package com.example.dynamicjavacode.util;


import com.example.dynamicjavacode.gitee.MemoryClassLoader;
import com.sun.tools.javac.file.JavacFileManager;
import org.springframework.boot.loader.jar.JarFile;

import javax.tools.*;
import java.io.*;
import java.net.URI;
import java.net.URL;
import java.nio.CharBuffer;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.stream.Collectors;

/**
 * 内存Java文件管理器
 * 用于加载springboot boot info lib 下面的依赖资源
 */
public class MemoryJavaFileManager extends ForwardingJavaFileManager<JavaFileManager> {

    // compiled classes in bytes:
    final Map<String, byte[]> classBytes = new HashMap<String, byte[]>();

    final Map<String, List<JavaFileObject>> classObjectPackageMap = new HashMap<>();

    private JavacFileManager javaFileManager;

    /**
     * key 包名 value javaobj 主要给jdk编译class的时候找依赖class用
     */
    public final static Map<String, List<JavaFileObject>> CLASS_OBJECT_PACKAGE_MAP = new HashMap<>();

    private static final Object lock = new Object();

    private static boolean isInit = false;


    public void init() {
        try {
            String jarBaseFile = MemoryClassLoader.getPath();
            JarFile jarFile = new JarFile(new File(jarBaseFile));
            List<JarEntry> entries = jarFile.stream().filter(jarEntry -> {
                return jarEntry.getName().endsWith(".jar");
            }).collect(Collectors.toList());


            JarFile libTempJarFile = null;
            List<JavaFileObject> onePackgeJavaFiles = null;
            String packgeName = null;

            List<JarEntry> projectCollect = jarFile.stream().filter(jarEntry -> {
                return jarEntry.getName().indexOf("BOOT-INF") != -1 && jarEntry.getName().endsWith(".class");
            }).collect(Collectors.toList());
            for (JarEntry entry : projectCollect) {
                String classPath = entry.getName().substring("BOOT-INF/classes/".length()).replace("/", ".").replace(".class", "");
                packgeName = classPath.substring(0, classPath.lastIndexOf("."));
                onePackgeJavaFiles = CLASS_OBJECT_PACKAGE_MAP.containsKey(packgeName) ? CLASS_OBJECT_PACKAGE_MAP.get(packgeName) : new ArrayList<>();
                onePackgeJavaFiles.add(new MemorySpringBootInfoJavaClassObject(entry.getName().substring("BOOT-INF/classes/".length()).replace("/", ".").replace(".class", ""),
                        new URL(jarFile.getUrl(), entry.getName()), javaFileManager));
                CLASS_OBJECT_PACKAGE_MAP.put(packgeName, onePackgeJavaFiles);
            }


            for (JarEntry entry : entries) {
                libTempJarFile = jarFile.getNestedJarFile(jarFile.getEntry(entry.getName()));
                if (libTempJarFile.getName().contains("tools.jar")) {
                    continue;
                }
                Enumeration<JarEntry> tempEntriesEnum = libTempJarFile.entries();
                while (tempEntriesEnum.hasMoreElements()) {
                    JarEntry jarEntry = tempEntriesEnum.nextElement();
                    String classPath = jarEntry.getName().replace("/", ".");
                    if (!classPath.endsWith(".class") || jarEntry.getName().lastIndexOf("/") == -1) {
                        continue;
                    } else {
                        packgeName = classPath.substring(0, jarEntry.getName().lastIndexOf("/"));
                        onePackgeJavaFiles = CLASS_OBJECT_PACKAGE_MAP.containsKey(packgeName) ? CLASS_OBJECT_PACKAGE_MAP.get(packgeName) : new ArrayList<>();
                        onePackgeJavaFiles.add(new MemorySpringBootInfoJavaClassObject(jarEntry.getName().replace("/", ".").replace(".class", ""),
                                new URL(libTempJarFile.getUrl(), jarEntry.getName()), javaFileManager));
                        CLASS_OBJECT_PACKAGE_MAP.put(packgeName, onePackgeJavaFiles);
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        isInit = true;

    }


    public MemoryJavaFileManager(JavaFileManager fileManager) {
        super(fileManager);
        this.javaFileManager = (JavacFileManager) fileManager;
    }

    public Map<String, byte[]> getClassBytes() {
        return new HashMap<String, byte[]>(this.classBytes);
    }

    @Override
    public void flush() throws IOException {
    }

    @Override
    public void close() throws IOException {
        classBytes.clear();
    }


    public List<JavaFileObject> getLibJarsOptions(String packgeName) {
        synchronized (lock) {
            if (!isInit) {
                init();
            }
        }
        return CLASS_OBJECT_PACKAGE_MAP.get(packgeName);
    }

    @Override
    public Iterable<JavaFileObject> list(Location location,
                                         String packageName,
                                         Set<JavaFileObject.Kind> kinds,
                                         boolean recurse)
            throws IOException {


        if ("CLASS_PATH".equals(location.getName()) && MemoryClassLoader.isJar()) {
            List<JavaFileObject> result = getLibJarsOptions(packageName);
            if (result != null) {
                return result;
            }
        }

        Iterable<JavaFileObject> it = super.list(location, packageName, kinds, recurse);

        if (kinds.contains(JavaFileObject.Kind.CLASS)) {
            final List<JavaFileObject> javaFileObjectList = classObjectPackageMap.get(packageName);
            if (javaFileObjectList != null) {
                if (it != null) {
                    for (JavaFileObject javaFileObject : it) {
                        javaFileObjectList.add(javaFileObject);
                    }
                }
                return javaFileObjectList;
            } else {
                return it;
            }
        } else {
            return it;
        }
    }

    @Override
    public String inferBinaryName(Location location, JavaFileObject file) {
        if (file instanceof MemoryInputJavaClassObject) {
            return ((MemoryInputJavaClassObject) file).inferBinaryName();
        }
        return super.inferBinaryName(location, file);
    }

    @Override
    public JavaFileObject getJavaFileForOutput(Location location, String className, JavaFileObject.Kind kind,
                                               FileObject sibling) throws IOException {
        if (kind == JavaFileObject.Kind.CLASS) {
            return new MemoryOutputJavaClassObject(className);
        } else {
            return super.getJavaFileForOutput(location, className, kind, sibling);
        }
    }

    public JavaFileObject makeStringSource(String className, final String code) {
        String classPath = className.replace('.', '/') + JavaFileObject.Kind.SOURCE.extension;

        return new SimpleJavaFileObject(URI.create("string:///" + classPath), JavaFileObject.Kind.SOURCE) {
            @Override
            public CharBuffer getCharContent(boolean ignoreEncodingErrors) {
                return CharBuffer.wrap(code);
            }
        };
    }

    void makeBinaryClass(String className, final byte[] bs) {
        JavaFileObject javaFileObject = new MemoryInputJavaClassObject(className, bs);

        String packageName = "";
        int pos = className.lastIndexOf('.');
        if (pos > 0) {
            packageName = className.substring(0, pos);
        }
        List<JavaFileObject> javaFileObjectList = classObjectPackageMap.get(packageName);
        if (javaFileObjectList == null) {
            javaFileObjectList = new LinkedList<>();
            javaFileObjectList.add(javaFileObject);

            classObjectPackageMap.put(packageName, javaFileObjectList);
        } else {
            javaFileObjectList.add(javaFileObject);
        }
    }

    public class MemoryInputJavaClassObject extends SimpleJavaFileObject {
        public final String className;
        public final byte[] bs;

        public MemoryInputJavaClassObject(String className, byte[] bs) {
            super(URI.create("string:///" + className.replace('.', '/') + Kind.CLASS.extension), Kind.CLASS);
            this.className = className;
            this.bs = bs;
        }

        @Override
        public InputStream openInputStream() {
            return new ByteArrayInputStream(bs);
        }

        public String inferBinaryName() {
            return className;
        }
    }


    public class MemoryOutputJavaClassObject extends SimpleJavaFileObject {
        public final String className;

        public MemoryOutputJavaClassObject(String className) {
            super(URI.create("string:///" + className.replace('.', '/') + Kind.CLASS.extension), Kind.CLASS);
            this.className = className;
        }

        @Override
        public OutputStream openOutputStream() {
            return new FilterOutputStream(new ByteArrayOutputStream()) {
                @Override
                public void close() throws IOException {
                    out.close();
                    ByteArrayOutputStream bos = (ByteArrayOutputStream) out;
                    byte[] bs = bos.toByteArray();
                    classBytes.put(className, bs);
                    makeBinaryClass(className, bs);
                }
            };
        }
    }
}
